import sys

# import fitz  # PyMuPDF
import pymupdf as fitz

from PySide6.QtWidgets import (
    QApplication,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QLabel,
    QFileDialog,
    QLineEdit,
    QMessageBox,
)
from PySide6.QtGui import QPixmap, QImage
from PySide6.QtCore import Qt
import os
from PySide6.QtGui import QShortcut, QKeySequence


class PDFTool(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("PDF工具")
        self.pdf_path = None
        self.doc = None

        layout = QVBoxLayout()

        # 选择PDF按钮
        self.select_btn = QPushButton("选择PDF文档")
        self.select_btn.clicked.connect(self.select_pdf)
        layout.addWidget(self.select_btn)

        # # PDF第一页上半部分缩略图
        # self.thumb_label = QLabel("缩略图预览")
        # self.thumb_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        # layout.addWidget(self.thumb_label)
        # 缩略图预览功能已废弃，仅显示文件名
        self.thumb_label = QLabel("未选择文件")
        self.thumb_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.thumb_label)

        # 逐页转为PNG按钮
        self.export_png_btn = QPushButton("逐页转为PNG")
        self.export_png_btn.clicked.connect(self.export_pngs)
        layout.addWidget(self.export_png_btn)

        # 页码输入和提取按钮
        hbox = QHBoxLayout()
        self.page_input = QLineEdit()
        self.page_input.setPlaceholderText("输入页码(从1开始)")
        hbox.addWidget(self.page_input)
        self.extract_btn = QPushButton("提取")
        self.extract_btn.clicked.connect(self.extract_page)
        hbox.addWidget(self.extract_btn)
        layout.addLayout(hbox)

        self.setLayout(layout)

        # 添加快捷键Esc退出程序
        QShortcut(QKeySequence("Esc"), self, self.close)

    def select_pdf(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择PDF文件", "", "PDF Files (*.pdf)"
        )
        if file_path:
            self.pdf_path = file_path
            self.doc = fitz.open(file_path)
            # 缩略图预览功能已废弃，仅显示文件名
            # self.show_thumbnail()
            file_name = os.path.basename(file_path)
            self.thumb_label.setText(f"已选择文件: {file_name}")

    @DeprecationWarning
    def show_thumbnail(self):
        if not self.doc or self.doc.page_count == 0:
            self.thumb_label.setText("无法加载PDF")
            return
        # page = self.doc.load_page(0)
        page = self.doc[0]
        pix = page.get_pixmap()  # type: ignore
        raw_pix = pix.tobytes()  # output="raw", alpha=False
        # 获取位图的宽度和高度
        # pix.save("page_image.png")
        # 只取上半部分
        # 选择合适的QImage格式
        # if pix.alpha:
        #     img_format = QImage.Format.Format_ARGB32
        # else:
        img_format = QImage.Format.Format_RGB32
        img = QImage(raw_pix, pix.width, pix.height, pix.stride, img_format)
        half_height = img.height() // 2
        cropped = img.copy(0, 0, img.width(), half_height)
        pixmap = QPixmap.fromImage(cropped).scaled(
            400, 200, Qt.AspectRatioMode.KeepAspectRatio
        )
        self.thumb_label.setPixmap(pixmap)

    def export_pngs(self):
        if not self.doc:
            QMessageBox.warning(self, "错误", "请先选择PDF文件")
            return
        out_dir = QFileDialog.getExistingDirectory(self, "选择输出文件夹")
        if not out_dir:
            return
        for i in range(self.doc.page_count):
            page = self.doc.load_page(i)
            pix = page.get_pixmap(dpi=300)  # type: ignore
            img_path = os.path.join(out_dir, f"page_{i+1}.png")
            pix.save(img_path)
        QMessageBox.information(self, "完成", "所有页面已导出为PNG")

    def extract_page(self):
        if not self.doc:
            QMessageBox.warning(self, "错误", "请先选择PDF文件")
            return
        try:
            page_num = int(self.page_input.text()) - 1
            if not (0 <= page_num < self.doc.page_count):
                raise ValueError
        except ValueError:
            QMessageBox.warning(self, "错误", "请输入有效的页码")
            return
        out_path, _ = QFileDialog.getSaveFileName(
            self, "保存提取页", f"page_{page_num+1}.pdf", "PDF Files (*.pdf)"
        )
        if not out_path:
            return
        new_doc = fitz.open()
        new_doc.insert_pdf(self.doc, from_page=page_num, to_page=page_num)
        new_doc.save(out_path)
        new_doc.close()
        QMessageBox.information(self, "完成", f"第{page_num+1}页已提取为PDF")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = PDFTool()
    win.resize(500, 400)
    win.show()
    sys.exit(app.exec())
